import copy


def allocate_replenishments_future_only(product_plan: dict, replenishment: dict) -> dict:
    updated = copy.deepcopy(product_plan)
    week_stock = {wk: {mat: int(q) for mat, q in replenishment[wk].items()} for wk in replenishment}

    # 周次转换函数
    def w2i(w):
        try:
            return int(w)
        except:
            return w

    # 按周次优先再产品顺序排序
    product_order = list(product_plan.keys())
    for pid in updated:
        weeks = list(updated[pid].keys())
        # 排除非周次字段
        week_keys = [w for w in weeks if w.isdigit() or w2i(w)]
        updated[pid]['_sorted_weeks'] = sorted(week_keys, key=w2i)

    # 按周次排序的产品列表
    pids_sorted = sorted(product_order, key=lambda pid: (
        min(w2i(w) for w in updated[pid]['_sorted_weeks']),  # 最早周次
        product_order.index(pid)
    ))

    # 主处理逻辑
    for pid in pids_sorted:
        pdata = updated[pid]
        for week in pdata['_sorted_weeks']:
            week_data = pdata[week]
            pweek = w2i(week)
            baseline = int(week_data.get('Baseline Demand', 0))
            fullkit = int(week_data.get('Full kit Num', 0))

            # 初始化补产结果
            week_data['Adjusted Full kit Num'] = fullkit
            week_data['Adjusted Full kit Rate'] = round(fullkit / baseline, 3) if baseline > 0 else 0

            if fullkit >= baseline:
                continue

            # 获取短缺信息
            shortages = copy.deepcopy(week_data.get('shortage_info', []))
            if not shortages:
                continue

            # 筛选可用补货周
            future_weeks = sorted([wk for wk in week_stock.keys() if w2i(wk) > pweek], key=w2i)
            last_update_week = None

            # 补货处理
            for rw in future_weeks:
                per_material_units = []
                for s in shortages:
                    mat = s['material']
                    unit_need = int(s['unit_need'])
                    shortage_qty = int(s['shortage'])
                    inv_qty = week_stock[rw].get(mat, 0)

                    if unit_need <= 0:
                        per_material_units.append(0)
                        continue
                    per_material_units.append(min(inv_qty // unit_need, shortage_qty // unit_need))

                if not per_material_units:
                    continue

                add_units = min(per_material_units)
                slack = baseline - fullkit
                add_units = min(add_units, slack)

                if add_units <= 0:
                    continue

                # 更新库存和短缺信息
                for s in shortages:
                    mat = s['material']
                    unit_need = int(s['unit_need'])
                    use_qty = add_units * unit_need
                    s['shortage'] = max(0, int(s['shortage']) - use_qty)
                    week_stock[rw][mat] = max(0, week_stock[rw].get(mat, 0) - use_qty)

                fullkit += add_units
                week_data['Adjusted Full kit Num'] = fullkit
                week_data['Adjusted Full kit Rate'] = round(fullkit / baseline, 3) if baseline > 0 else 0
                last_update_week = rw

                if fullkit >= baseline:
                    break

            if last_update_week:
                week_data['update_from'] = last_update_week

    return updated


if __name__ == '__main__':
    product_plan = {
        "P1": {
            "202534": {
                "Baseline Demand": 200,
                "Full kit Num": 100,
                "Full kit Rate": 0.5,
                "shortage_info": [
                    {"material": "M1", "unit_need": 1, "shortage": 100},
                    {"material": "M2", "unit_need": 2, "shortage": 200}
                ]
            },
            "202535": {
                "Baseline Demand": 150,
                "Full kit Num": 80,
                "Full kit Rate": 0.533,
                "shortage_info": [
                    {"material": "M1", "unit_need": 1, "shortage": 70},
                    {"material": "M3", "unit_need": 1, "shortage": 70}
                ]
            }
        },
        "P2": {
            "202534": {
                "Baseline Demand": 200,
                "Full kit Num": 100,
                "Full kit Rate": 0.5,
                "shortage_info": [
                    {"material": "M1", "unit_need": 2, "shortage": 200},
                    {"material": "M3", "unit_need": 2, "shortage": 200}
                ]
            }
        }
    }

    replenishment = {
        "202535": {"M1": 200, "M2": 300, "M3": 200, "M4": 200},
        "202536": {"M1": 150, "M2": 50, "M3": 100}
    }

    updated = allocate_replenishments_future_only(product_plan, replenishment)
    import json
    print(json.dumps(updated, indent=2, ensure_ascii=False))
